/*
 * CudaAESInvOperations.cpp
 *
 *  Created on: Feb 26, 2010
 *      Author: Łukasz Spintzyk
 */

#include <cuda.h>
#include "aesUtils.h"
#include "CudaAESInvOperations.h"
#include "CudaAESOperations.h"
#include "cudaUtils.h"

__device__ void cudaInvByteSub(byte* state)
{
    int col = threadIdx.x;
    int row = threadIdx.y;

    byte* cell = ((byte*) state + col * 4) + row;
    *cell = DSBOX_INV[*cell];
}

__global__ void KcudaInvByteSub(cudaPitchedPtr state)
{
    extern __shared__ byte shared[];
    byte* data2 = shared;

    
    int col = threadIdx.x;
    int row = threadIdx.y;
    
    byte* destCell = ((byte*) state.ptr + col * state.pitch)
            + row;
    byte* p = data2 + col*4 + row;
    *p = *destCell;
    
    cudaInvByteSub(data2);
    
    *destCell = *p;
}

__device__ void cudaInvShiftRow(byte* state, byte* outputState, int nb)
{
    int col = threadIdx.x;
    int row = threadIdx.y;


    int dest_col =  DSHIFT_NB[nb][row][col];
    byte* destCell = ((byte*) outputState + col * 4) + row;
    
    byte* srcCell = ((byte*) state + dest_col*4) + row;
    
    *destCell = *srcCell;
}

__global__ void KcudaInvShiftRow(cudaPitchedPtr state, cudaPitchedPtr outputState,
        int nb)
{
    extern __shared__ byte shared[];
    byte* data2 = shared;
    byte* buffer2 = shared+32;
    
    int col = threadIdx.x;
    int row = threadIdx.y;
    
    byte* srcCell = ((byte*) state.ptr + col * state.pitch)
            + row;
    byte* destCell = ((byte*) outputState.ptr + col * outputState.pitch) + row;
    
    byte* p = data2 + col*4 + row;
    *p = *srcCell;
    byte* pb = buffer2 + col*4 + row;
    
    cudaInvShiftRow(data2, buffer2, nb);
    *destCell = *pb;
    
}


__device__ void cudaInvMixColumn(byte* state, byte* outputState, int nb)
{
    int col = threadIdx.x;
    int row = threadIdx.y;

    byte* destCell = ((byte*) outputState + col * 4) + row;
    byte* srcCell = ((byte*) state + col * 4);
    
    byte* a0 = srcCell +row;
    byte* a1 = srcCell + ((1 + row)&0x3);
    byte* a2 = srcCell + ((2 + row)&0x3);
    byte* a3 = srcCell + ((3 + row)&0x3);

    byte x1, x2, x4, x8;
    
    x2 = XTIME(*a0);
    x4 = XTIME(x2);
    x8 = XTIME(x4);
    *destCell = x2 ^ x4 ^x8;
    
    x1 = *a1;
    x2 = XTIME(x1);
    x8 = XTIME(XTIME(x2));
    *destCell ^= x1^x2^x8;
    
    x1 = *a2;
    x4 = XTIME(XTIME(x1));
    x8 = XTIME(x4);
    *destCell ^= x1^x4^x8;
  
    x1 = *a3;
    x8 = XTIME(XTIME(XTIME(x1)));
    *destCell ^= x1^x8;
    
}
 
__global__ void KcudaInvMixColumn(cudaPitchedPtr state, cudaPitchedPtr outputState,
        int nb)
{
    extern __shared__ byte shared[];
    byte* data2 = shared;
    byte* buffer2 = shared+32;
    
    int col = threadIdx.x;
    int row = threadIdx.y;
    
    byte* srcCell = ((byte*) state.ptr + col * state.pitch)
            + row;
    byte* destCell = ((byte*) outputState.ptr + col * outputState.pitch) + row;
    
    byte* p = data2 + col*4 + row;
    *p = *srcCell;
    byte* pb = buffer2 + col*4 + row;
    
    cudaInvMixColumn(data2, buffer2, nb);
    
    *destCell = *pb;
}

__global__ void KcudaInvExpandedKey(cudaPitchedPtr eKey, cudaPitchedPtr workBuffer, int nb, int nr)
{
    extern __shared__ byte shared[];
    byte* data2 = shared;
    byte* buffer2 = shared+32;
    
    int col = threadIdx.x;
    int row = threadIdx.y;

    unsigned int wb = BLOCK_SIZE(nb)/4;
    byte* destCell = ((byte*) eKey.ptr + col * eKey.pitch) + row;
   // byte* srcCell = ((byte*) workBuffer.ptr + col * workBuffer.pitch) + row;
    
    
    cudaPitchedPtr tkey;
    tkey.ptr = eKey.ptr;
    tkey.pitch = eKey.pitch;
    tkey.xsize = eKey.xsize;
    tkey.ysize = eKey.ysize;
    tkey.ptr = (void**)((byte*)eKey.ptr +(wb*eKey.pitch));
    
    
    for(int i = 1; i < nr; ++i)
    {
        tkey.ptr = (void**)((byte*)eKey.ptr +((i*wb + col)*eKey.pitch) + row);
        //destCell = ((byte*) tkey.ptr + ((i*wb)*tkey.pitch));// + row;
        destCell = ((byte*) eKey.ptr + ((col + i*wb) * eKey.pitch)) + row;
        
        byte* p = data2 + col*4 + row;
        *p = *((byte*)tkey.ptr);
        
        byte* pb = buffer2 + col*4+row;

        cudaInvMixColumn(data2, buffer2, nb);
        
        
        *destCell=*pb;
    }
    
}

__device__ void cudaAesInvRound(byte* state, byte* cState, cudaPitchedPtr invEKey, int nb)
{
/*This is other method of decrypting 
    cudaAddRoundKey(state, invEKey, nb);
    cudaInvMixColumn(state, cState, nb);
    cudaInvShiftRow(cState, state, nb);
    cudaInvByteSub(state);
*/
    cudaInvByteSub(state);
    // __syncthreads();
    cudaInvShiftRow(state, cState, nb);
    //__syncthreads();
    // now in cState is right State
    cudaInvMixColumn(cState, state, nb);
    // __syncthreads();
    //now in state is right state
    cudaAddRoundKey(state, invEKey, nb);
    //cudaCopyPitchedPtr(state, eKey);
}

__global__ void KcudaAesInvRound(cudaPitchedPtr state, cudaPitchedPtr cState, cudaPitchedPtr eKey, int nb)
{
    extern __shared__ byte shared[];
    byte* data2 = shared;
    byte* buffer2 = shared+32;
    
    int col = threadIdx.x;
    int row = threadIdx.y;
    
    byte* srcCell = ((byte*) state.ptr + col * state.pitch)
            + row;
    byte* destCell = ((byte*) cState.ptr + col * cState.pitch) + row;
    
    byte* p = data2 + col*4 + row;
    *p = *srcCell;
    byte* pb = buffer2 + col*4 + row;
    
    
    cudaAesInvRound(data2, buffer2, eKey, nb);
    
    *destCell = *pb;
    
}

__device__ void cudaAesInvFinalRound(byte* state, byte* outState, cudaPitchedPtr invEKey, int nb)
{
/*  This is other method of decrypting 
    cudaAddRoundKey(state, invEKey, nb);
    cudaInvShiftRow(state, outState, nb);
    cudaInvByteSub(outState);
*/
    cudaInvByteSub(state);
    // __syncthreads();
    cudaInvShiftRow(state, outState, nb);
    //__syncthreads();
    // now in outState is right State
    cudaAddRoundKey(outState, invEKey, nb);
}

__device__ void cudaInvRijndael(byte* state, byte* outState, cudaPitchedPtr invEKey, int nb, int nr)
{
/* This is other method of decrypting 
    unsigned int wb = BLOCK_SIZE(nb)/4;
    
    
    
    cudaPitchedPtr tkey;
    tkey.ptr = invEKey.ptr;
    tkey.pitch = invEKey.pitch;
    tkey.xsize = invEKey.xsize;
    tkey.ysize = invEKey.ysize;
    tkey.ptr = (void**)((byte*)invEKey.ptr +(nr*wb*invEKey.pitch));
    cudaAesInvFinalRound(state, outState, tkey, nb);
    
    for(int i = nr-1; i > 0; i--)
    {
        tkey.ptr = (void**)((byte*)invEKey.ptr +(i*wb*invEKey.pitch));
        cudaAesInvRound(outState, state, tkey, nb);
    }
    
    cudaAddRoundKey(outState, invEKey, nb);
*/
    
    unsigned int wb = BLOCK_SIZE(nb)/4;
    cudaAddRoundKey(state, invEKey, nb);
    
    cudaPitchedPtr tkey;
    tkey.ptr = invEKey.ptr;
    tkey.pitch = invEKey.pitch;
    tkey.xsize = invEKey.xsize;
    tkey.ysize = invEKey.ysize;
    tkey.ptr = (void**)((byte*)invEKey.ptr +(wb*invEKey.pitch));
    
    
    for(int i = 1; i < nr; ++i)
    {
        tkey.ptr = (void**)((byte*)invEKey.ptr +(i*wb*invEKey.pitch));
        cudaAesInvRound(state, outState, tkey, nb);
    }
    
    tkey.ptr = (void**)((byte*)invEKey.ptr +(nr*wb*invEKey.pitch));
    cudaAesInvFinalRound(state, outState, tkey, nb);
    
}


__global__ void KcudaInvRijndael(cudaPitchedPtr state, cudaPitchedPtr outState, cudaPitchedPtr invEKey, int nb, int nr)
{
    extern __shared__ byte shared[];
    byte* data2 = shared;
    byte* buffer2 = shared+32;
    
    int col = threadIdx.x;
    int row = threadIdx.y;
    
    byte* srcCell = ((byte*) state.ptr + col * state.pitch)
            + row;
    byte* destCell = ((byte*) outState.ptr + col * outState.pitch) + row;
    
    byte* p = data2 + col*4 + row;
    *p = *srcCell;
    byte* pb = buffer2 + col*4 + row;
    
    
    cudaInvRijndael(data2, buffer2, invEKey, nb, nr);
    
    *destCell = *pb;
}

void invByteSub(byte* state, const NB nb)
{    
    cudaPitchedPtr h_state;
    unsigned int size = BLOCK_SIZE(nb);
    unsigned int width = size / 4;
    unsigned int height = 4;
    
    h_state.xsize = height;
    h_state.ysize = width;
    
    CUDA_CHECK(cudaMallocPitch((void**) &h_state.ptr, &h_state.pitch,
            h_state.xsize, h_state.ysize));
    
    CUDA_CHECK(cudaMemcpy2D(h_state.ptr, h_state.pitch, state, height, height,
            width, cudaMemcpyHostToDevice));
    
    dim3 dimBlock(width, height);
    KcudaInvByteSub<<<1, dimBlock>>>(h_state);

    CUDA_CHECK(cudaMemcpy2D(state, height, h_state.ptr, h_state.pitch, height,
            width, cudaMemcpyDeviceToHost));
    
    CUDA_CHECK(cudaFree(h_state.ptr));
}

void invShiftRow(byte* state, const NB nb)
{
    assert(nb <= NB_8);
    unsigned int size = BLOCK_SIZE(nb);
    unsigned int width = size / 4;
    unsigned int height = 4;
    
    cudaPitchedPtr d_state;
    cudaPitchedPtr d_outputState;
    
    d_state.xsize = height;
    d_state.ysize = width;
    
    d_outputState.xsize = height;
    d_outputState.ysize = width;
    
    CUDA_CHECK(cudaMallocPitch((void**) &d_state.ptr, &d_state.pitch,
            d_state.xsize, d_state.ysize));
    CUDA_CHECK(cudaMallocPitch((void**) &d_outputState.ptr,
            &d_outputState.pitch, d_outputState.xsize, d_outputState.ysize));
    
    CUDA_CHECK(cudaMemcpy2D(d_state.ptr, d_state.pitch, state, height, height,
            width, cudaMemcpyHostToDevice));
    
    dim3 dimBlock(width, height);//(w, h)
    KcudaInvShiftRow<<<1, dimBlock>>>(d_state, d_outputState, nb);

    CUDA_CHECK(cudaMemcpy2D(state, height, d_outputState.ptr,
            d_outputState.pitch, height, width, cudaMemcpyDeviceToHost));
    
    CUDA_CHECK(cudaFree(d_state.ptr));
    CUDA_CHECK(cudaFree(d_outputState.ptr));
}

void invMixColumn(byte* state, const NB nb)
{
    assert(nb <= NB_8);
    unsigned int size = BLOCK_SIZE(nb);
    unsigned int width = size / 4;
    unsigned int height = 4;
    
    cudaPitchedPtr d_state;
    cudaPitchedPtr d_outputState;
    
    d_state.xsize = height;
    d_state.ysize = width;
    
    d_outputState.xsize = height;
    d_outputState.ysize = width;
    
    CUDA_CHECK(cudaMallocPitch((void**) &d_state.ptr, &d_state.pitch,
            d_state.xsize, d_state.ysize));
    CUDA_CHECK(cudaMallocPitch((void**) &d_outputState.ptr,
            &d_outputState.pitch, d_outputState.xsize, d_outputState.ysize));
    
    CUDA_CHECK(cudaMemcpy2D(d_state.ptr, d_state.pitch, state, height, height,
            width, cudaMemcpyHostToDevice));
    
    dim3 dimBlock(width, height);//(w, h)
    KcudaInvMixColumn<<<1, dimBlock>>>(d_state, d_outputState, nb);

    CUDA_CHECK(cudaMemcpy2D(state, height, d_outputState.ptr,
            d_outputState.pitch, height, width, cudaMemcpyDeviceToHost));
    
    CUDA_CHECK(cudaFree(d_state.ptr));
    CUDA_CHECK(cudaFree(d_outputState.ptr));
}


void invExpandedKey(byte* eKey, const NB nb, const NB nk)
{
    assert(nb <= NB_8);
    unsigned int size = BLOCK_SIZE(nb);
    unsigned int width = size / 4;
    unsigned int height = 4;
    
    cudaPitchedPtr d_expKey;
    cudaPitchedPtr d_buffer;
    
    d_expKey.xsize = height;
    d_expKey.ysize = NR[nk][nb]*width*4;
    
    d_buffer.xsize = height;
    d_buffer.ysize = width;
    
    CUDA_CHECK(cudaMallocPitch((void**) &d_expKey.ptr, &d_expKey.pitch,
            d_expKey.xsize, d_expKey.ysize));
    CUDA_CHECK(cudaMallocPitch((void**) &d_buffer.ptr,
            &d_buffer.pitch, d_buffer.xsize, d_buffer.ysize));
    
    CUDA_CHECK(cudaMemcpy2D(d_expKey.ptr, d_expKey.pitch, eKey, height, height,
            NR[nk][nb]*width*4, cudaMemcpyHostToDevice));
    
    dim3 dimBlock(width, height);//(w, h)
    KcudaInvExpandedKey<<<1, dimBlock>>>(d_expKey, d_buffer, nb, NR[nk][nb]);

    CUDA_CHECK(cudaMemcpy2D(eKey, height, d_expKey.ptr,
            d_expKey.pitch, height, d_expKey.ysize, cudaMemcpyDeviceToHost));
    
    CUDA_CHECK(cudaFree(d_expKey.ptr));
    CUDA_CHECK(cudaFree(d_buffer.ptr));
}

void aesDecryptInit()
{
//    cuInit(0);
    CUDA_CHECK(cudaMemcpyToSymbol(DSBOX_INV, SBOX_INV, 256*sizeof(byte), 0, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpyToSymbol(DSHIFT_NB, SHIFT_NB, SHIFT_NB_SIZE, 0, cudaMemcpyHostToDevice));
}


void aesInvRound(byte* state, byte* eKey, const NB nb)
{
    assert(nb <= NB_8);
    unsigned int size = BLOCK_SIZE(nb);
    unsigned int width = size / 4;
    unsigned int height = 4;
    
    cudaPitchedPtr d_state;
    cudaPitchedPtr d_outputState;
    cudaPitchedPtr d_key;
    
    d_state.xsize = height;
    d_state.ysize = width;
    
    d_outputState.xsize = height;
    d_outputState.ysize = width;
    
    d_key.xsize = height;
    d_key.ysize = width;
    
    CUDA_CHECK(cudaMallocPitch((void**) &d_state.ptr, &d_state.pitch,
            d_state.xsize, d_state.ysize));
    
    CUDA_CHECK(cudaMallocPitch((void**) &d_outputState.ptr,
            &d_outputState.pitch, d_outputState.xsize, d_outputState.ysize));
    
    CUDA_CHECK(cudaMallocPitch((void**) &d_key.ptr,
            &d_key.pitch, d_key.xsize, d_key.ysize));
    
    
    CUDA_CHECK(cudaMemcpy2D(d_state.ptr, d_state.pitch, state, height, height,
            width, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy2D(d_key.ptr, d_key.pitch, eKey, height, height,
            width, cudaMemcpyHostToDevice));

    
    dim3 dimBlock(width, height);//(w, h)
    KcudaAesInvRound<<<1, dimBlock>>>(d_state, d_outputState, d_key, nb);
    
    
    CUDA_CHECK(cudaMemcpy2D(state, height, d_state.ptr,
            d_state.pitch, height, width, cudaMemcpyDeviceToHost));
    
    CUDA_CHECK(cudaFree(d_state.ptr));
    CUDA_CHECK(cudaFree(d_outputState.ptr));
    CUDA_CHECK(cudaFree(d_key.ptr));
}

void aesDecryptBlock(byte* state, byte* eKey, const NB nb, const NB nk)
{
    assert(nb <= NB_8);
    unsigned int size = BLOCK_SIZE(nb);
    unsigned int width = size / 4;
    unsigned int height = 4;
    
    cudaPitchedPtr d_state;
    cudaPitchedPtr d_outputState;
    cudaPitchedPtr d_key;
    
    d_state.xsize = height;
    d_state.ysize = width;
    
    d_outputState.xsize = height;
    d_outputState.ysize = width;
    
    d_key.xsize = height;
    d_key.ysize = NR[nk][nb]*width*4;// todo check if it is good size
    
    CUDA_CHECK(cudaMallocPitch((void**) &d_state.ptr, &d_state.pitch,
            d_state.xsize, d_state.ysize));
    
    CUDA_CHECK(cudaMallocPitch((void**) &d_outputState.ptr,
            &d_outputState.pitch, d_outputState.xsize, d_outputState.ysize));
    
    CUDA_CHECK(cudaMallocPitch((void**) &d_key.ptr,
            &d_key.pitch, d_key.xsize, d_key.ysize));
    
    
    CUDA_CHECK(cudaMemcpy2D(d_state.ptr, d_state.pitch, state, height, height,
            width, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy2D(d_key.ptr, d_key.pitch, eKey, height, height,
            NR[nk][nb]*width*4, cudaMemcpyHostToDevice));

    unsigned int shMemSize = 32*2 + 15*32;
    dim3 dimBlock(width, height);//(w, h)
    KcudaInvRijndael<<<1, dimBlock, shMemSize>>>(d_state, d_outputState, d_key, nb, NR[nk][nb]);
    
    
    CUDA_CHECK(cudaMemcpy2D(state, height, d_outputState.ptr,
            d_outputState.pitch, height, width, cudaMemcpyDeviceToHost));
    
    CUDA_CHECK(cudaFree(d_state.ptr));
    CUDA_CHECK(cudaFree(d_outputState.ptr));
    CUDA_CHECK(cudaFree(d_key.ptr));
}
